#include <math.h>
#include <stdio.h>
#include <stdlib.h>
#include <sys/time.h>
#include <cuda_runtime.h>
#include <cuComplex.h>


#define SIGNAL_SIZE 2048
#define CHECK(call)\
{\
    const cudaError_t error = call;\
    if(error != cudaSuccess)\
    {\
        printf("Error: %s:%d", __FILE__, __LINE__);\
        printf("Code:%d, reason: %s\n",error, cudaGetErrorString(error));\
        exit(1);\
    }\
}

__global__ void FFT_transform(cuComplex *input,
                              cuComplex *input_temp,
                              cuComplex *Omega,
                              int size,
                              int scope,
                              int k);

__global__ void FFT_transform_coalesce(cuComplex *input,
                                       cuComplex *input_temp,
                                       cuComplex *Omega,
                                       int size,
                                       int scope,
                                       int k);

void reader_transform(cuComplex *x, int size);
void FFT_transform_host(int n, cuComplex *x, cuComplex *w);
void checkresult(cuComplex *x, cuComplex *y, int size);
void swap(cuComplex *x, cuComplex *y);
double cpusecond();

void reader_transform(cuComplex *x, int size){
  int i, j, l;
  for (i = 0, j = 0; i < size; ++i) {
      if (i < j)
          swap(&x[i], &x[j]);
      for (l = size >> 1; (j ^= l) < l; l >>=1);
  }
  return;
}

void FFT_transform_host(int n, cuComplex *x, cuComplex *w){
  for (int i = 2; i <= n; i <<= 1) {
   int m = i >> 1;
   for (int j = 0; j < n; j += i) {
    for (int k = 0; k != m; ++k) {
     cuComplex z = cuCmulf(x[j + m + k],w[n / i * k]);
     x[j + m + k] = cuCsubf(x[j + k],z);
     x[j + k] =cuCaddf(x[j+k],z);
    }
   }
  }
  return;
}

void checkresult(cuComplex *x, cuComplex *y, int size){
  int i;
  int error = 0;
  for (i = 0; i < size; ++i) {
    if (fabs(x[i].x-y[i].x) > 0.1)
      error++;
    if (fabs(x[i].y-y[i].y) > 0.1)
      error++;
  }
  if (error > 0)
    printf("result wrong\n");
  else
    printf("result right\n");
  return;
}

void swap(cuComplex *x, cuComplex *y){
  cuComplex temp;
  temp = *x;
  *x = *y;
  *y = temp;
  return;
}

double cpusecond(){
  struct timeval tp;
  gettimeofday(&tp, NULL);
  return ((double)tp.tv_sec + (double) tp.tv_usec*1.e-6);
}

__global__ void FFT_transform(cuComplex *input,
                              cuComplex *input_temp,
                              cuComplex *Omega,
                              int size,
                              int scope,
                              int k){
  int idx = blockIdx.x * blockDim.x + threadIdx.x;
  int offset = scope >> 1;
  int prefix = size >>(k+1);
  int my_partner;
  int omega_inpart;
  cuComplex temp;

  if (idx < size) {
    my_partner = idx ^ offset;
    omega_inpart = idx & ((1<<k)-1);
    if (idx & offset) {
      temp = cuCmulf(input[idx],Omega[prefix * omega_inpart]);
      input_temp[idx] = cuCsubf(input[my_partner],temp);
      //printf("a[%d] = a[%d] - w[%d]*a[%d]\n", idx,my_partner,(prefix * (idx % offset)),idx); 
    } else {
      temp = cuCmulf(input[my_partner],Omega[prefix * omega_inpart]);
      input_temp[idx] = cuCaddf(input[idx],temp);
      //printf("a[%d] = a[%d] + w[%d]*a[%d]\n", idx,idx,(prefix * (idx % offset)),my_partner);
    }
  }
}

__global__ void FFT_transform_coalesce(cuComplex *input,
                                       cuComplex *input_temp,
                                       cuComplex *Omega,
                                       int size,
                                       int scope,
                                       int k){
  int idx = blockIdx.x * blockDim.x + threadIdx.x;
  int offset = scope >> 1;
  int prefix = size >>(k+1);
  int my_partner;
  int omega_inpart;
  cuComplex temp, complex_myself, complex_myparter, omega;

  if (idx < size) {
    my_partner = idx ^ offset;
    omega_inpart = idx & ((1<<k)-1);
    complex_myself = input[idx];
    complex_myparter = input[my_partner];
    omega = Omega[prefix * omega_inpart];
    if (idx & offset) {
      temp = cuCmulf(complex_myself,omega);
      input_temp[idx] = cuCsubf(complex_myparter,temp);
      //printf("a[%d] = a[%d] - w[%d]*a[%d]\n", idx,my_partner,(prefix * (idx % offset)),idx); 
    } else {
      temp = cuCmulf(complex_myparter,omega);
      input_temp[idx] = cuCaddf(complex_myself,temp);
      //printf("a[%d] = a[%d] + w[%d]*a[%d]\n", idx,idx,(prefix * (idx % offset)),my_partner);
    }
  }
}

int main(int argc, char **argv){
  int i,k,power_count;
  double istart;
  double ielaps;
  cuComplex *h_signal;
  cuComplex *h_signal_result_gpu;
  cuComplex *h_signal_result_cpu;
  cuComplex *d_signal;
  cuComplex *d_signal_temp;
  cuComplex *h_omega;
  cuComplex *d_omega;
  
  //allocate host memory
  h_signal = (cuComplex *)malloc(sizeof(cuComplex) * SIGNAL_SIZE);
  h_signal_result_gpu = (cuComplex *)malloc(sizeof(cuComplex) * SIGNAL_SIZE);
  h_signal_result_cpu = (cuComplex *)malloc(sizeof(cuComplex) * SIGNAL_SIZE);
  h_omega = (cuComplex *)malloc(sizeof(cuComplex) * SIGNAL_SIZE);
  
  //allocate device memory
  CHECK(cudaMalloc((cuComplex **)&d_signal,sizeof(cuComplex) * SIGNAL_SIZE));
  CHECK(cudaMalloc((cuComplex **)&d_signal_temp,sizeof(cuComplex) * SIGNAL_SIZE));
  CHECK(cudaMalloc((cuComplex **)&d_omega,sizeof(cuComplex) * SIGNAL_SIZE));
  
  //initial input
  for (i = 0; i < SIGNAL_SIZE; ++i) {
    h_signal[i].x = float(i);
    h_signal[i].y = 0.0;
  }
  //initial omega     
  for (i = 0; i < SIGNAL_SIZE; ++i) {
    h_omega[i].x = float(cos(2 * M_PI * i / SIGNAL_SIZE));
    h_omega[i].y = float(sin(2 * M_PI * i / SIGNAL_SIZE));
  }
  
  reader_transform(h_signal,SIGNAL_SIZE);
  memcpy(h_signal_result_cpu,h_signal, sizeof(cuComplex)*SIGNAL_SIZE);

  istart = cpusecond();
  FFT_transform_host(SIGNAL_SIZE,h_signal_result_cpu,h_omega);
  ielaps = cpusecond() - istart;
  printf("cpu time is: %f\n",ielaps);

  dim3 block(128,1);
  dim3 grid((SIGNAL_SIZE+block.x-1)/block.x);
  /*
  istart = cpusecond();
  //copy h_signal to d_signal and copy h_omega to d_omega
  CHECK(cudaMemcpy(d_signal, h_signal, sizeof(cuComplex) * SIGNAL_SIZE, cudaMemcpyHostToDevice));
  CHECK(cudaMemcpy(d_omega, h_omega, sizeof(cuComplex) * SIGNAL_SIZE, cudaMemcpyHostToDevice));
  for (i = 2,k=0; i <= SIGNAL_SIZE; i <<= 1) {
    FFT_transform <<<grid,block>>>(d_signal,d_signal_temp,d_omega,SIGNAL_SIZE,i,k);
    CHECK(cudaMemcpy(d_signal,d_signal_temp,sizeof(cuComplex) * SIGNAL_SIZE, cudaMemcpyDeviceToDevice));
    k++;
  }

  CHECK(cudaMemcpy(h_signal_result_gpu, d_signal, sizeof(cuComplex) * SIGNAL_SIZE, cudaMemcpyDeviceToHost));
  ielaps = cpusecond() - istart;
  checkresult(h_signal_result_cpu,h_signal_result_gpu,SIGNAL_SIZE);
  printf("gpu time is: %f\n",ielaps);
  */
  istart = cpusecond();
  //copy h_signal to d_signal and copy h_omega to d_omega
  CHECK(cudaMemcpy(d_signal, h_signal, sizeof(cuComplex) * SIGNAL_SIZE, cudaMemcpyHostToDevice));
  CHECK(cudaMemcpy(d_omega, h_omega, sizeof(cuComplex) * SIGNAL_SIZE, cudaMemcpyHostToDevice));
  power_count = 0;
  for (i = 2; i <=SIGNAL_SIZE; i<<=1)
      power_count++;
  if ((power_count%2) == 0) {
      for (i = 2,k=0; i < SIGNAL_SIZE; i <<= 1) {
        FFT_transform_coalesce <<<grid,block>>>(d_signal,d_signal_temp,d_omega,SIGNAL_SIZE,i,k);
        k++;
        i <<= 1;
        FFT_transform_coalesce <<<grid,block>>>(d_signal_temp,d_signal,d_omega,SIGNAL_SIZE,i,k);
        k++;
      }
      CHECK(cudaMemcpy(h_signal_result_gpu, d_signal, sizeof(cuComplex) * SIGNAL_SIZE, cudaMemcpyDeviceToHost));
  } else {
      for (i = 2,k=0; i < SIGNAL_SIZE; i <<= 1) {
        FFT_transform_coalesce <<<grid,block>>>(d_signal,d_signal_temp,d_omega,SIGNAL_SIZE,i,k);
        k++;
        i <<= 1;
        FFT_transform_coalesce <<<grid,block>>>(d_signal_temp,d_signal,d_omega,SIGNAL_SIZE,i,k);
        k++;
      }
      FFT_transform_coalesce <<<grid,block>>>(d_signal,d_signal_temp,d_omega,SIGNAL_SIZE,i,k);
      CHECK(cudaMemcpy(h_signal_result_gpu, d_signal_temp, sizeof(cuComplex) * SIGNAL_SIZE, cudaMemcpyDeviceToHost));
  }
  ielaps = cpusecond() - istart;
  checkresult(h_signal_result_cpu,h_signal_result_gpu,SIGNAL_SIZE);
  printf("gpu time is: %f\n",ielaps);

  free(h_signal);
  free(h_omega);
  free(h_signal_result_cpu);
  free(h_signal_result_gpu);

  cudaFree(d_signal);
  cudaFree(d_signal_temp);
  cudaFree(d_omega);
  /*
  for(i=0;i<SIGNAL_SIZE;++i)
  {
      printf("i:%d real:%f image:%f\n",i,h_signal_result_cpu[i].x,h_signal_result_cpu[i].y);
  }
  printf("\n");
  for(i=0;i<SIGNAL_SIZE;++i)
  {
      printf("i:%d real:%f image:%f\n",i,h_signal_result_gpu[i].x,h_signal_result_gpu[i].y);
  }
  */
  //checkresult(h_signal_result_cpu,h_signal_result_gpu,SIGNAL_SIZE);
  return 0;
}